import torch


class MyNet(torch.nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.seq = torch.nn.Sequential(
            torch.nn.Conv2d(3, 32, 5, 1, 2),
            torch.nn.MaxPool2d(2),
            torch.nn.Conv2d(32, 32, 5, 1, 2),
            torch.nn.MaxPool2d(2),
            torch.nn.Conv2d(32, 64, 5, 1, 2),
            torch.nn.MaxPool2d(2),
            torch.nn.Flatten(),
            torch.nn.Linear(64 * 4 * 4, 64),
            torch.nn.Linear(64, 10)
        )

    def forward(self, data_):
        return self.seq(data_)


if __name__ == '__main__':
    net = MyNet()
    print(net)

    x = torch.Tensor(64, 3, 32, 32)
    y = net(x)
    print(y)
